package com.qupeng.demo.kafka.kafkaapache.consumer;

import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.WakeupException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.time.Instant;
import java.time.ZoneId;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

public class MyConsumer {

    private final static Logger logger = LoggerFactory.getLogger(MyConsumer.class);

    public KafkaConsumer create() {
        Properties props = new Properties();
        props.put("bootstrap.servers", "localhost:9092");
        props.put("group.id", "test-group");
        props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");

        return new KafkaConsumer<>(props);
    }

    public void consume() {

        KafkaConsumer<String, String> consumer = create();
        consumer.subscribe(Arrays.asList("topic1", "topic2"));

        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
            // 处理拉取到的消息
            records.forEach(record -> System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value()));
        }
    }

    public void commitOffset() {
        KafkaConsumer<String, String> consumer = create();
        consumer.subscribe(Arrays.asList("topic1", "topic2"));
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
            }
            try {
                consumer.commitSync();
            } catch (CommitFailedException | WakeupException e) {
                logger.error("Commit failed.", e);
            }
        }
    }

    public void commitOffsetAsync() {
        KafkaConsumer<String, String> consumer = create();
        consumer.subscribe(Arrays.asList("topic1", "topic2"));
        while (true) {
            ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
            for (ConsumerRecord<String, String> record : records) {
                System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
            }
            consumer.commitAsync(new OffsetCommitCallback() {
                @Override
                public void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
                    if (null != exception) {
                        logger.error("Commit failed for offsets {}", offsets, exception);
                    }
                }
            });
        }
    }

    public void commitOffsetAsyncAndSync() {
        KafkaConsumer<String, String> consumer = create();
        consumer.subscribe(Arrays.asList("topic1", "topic2"));
        try {
            while (true) {
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
                for (ConsumerRecord<String, String> record : records) {
                    System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
                }
                consumer.commitAsync();
            }
        } catch (WakeupException e) {
            logger.error("Unexpected error", e);
            consumer.commitSync();
        } finally {
            consumer.close();
        }
    }

    public void commitSpecificOffset() {
        KafkaConsumer<String, String> consumer = create();
        consumer.subscribe(Arrays.asList("topic1", "topic2"));

        try {
            while (true) {
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
                for (ConsumerRecord<String, String> record : records) {
                    System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
                }
                // 手动提交偏移量
                for (TopicPartition partition : records.partitions()) {
                    long offset = records.records(partition).get(records.records(partition).size() - 1).offset();
                    consumer.commitSync(Collections.singletonMap(partition, new org.apache.kafka.clients.consumer.OffsetAndMetadata(offset + 1)));
                }
            }
        } catch (WakeupException e) {
            // Ignore for shutdown
        } finally {
            consumer.close();
        }
    }

    public void commitInRebalance() {
        KafkaConsumer<String, String> consumer = create();
        Map<TopicPartition, OffsetAndMetadata> currentOffsets = new HashMap<>();
        try {
            consumer.subscribe(Arrays.asList("topic1", "topic2"), new MyRebalanceListener(consumer));
            while (true) {
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
                for (ConsumerRecord<String, String> record : records) {
                    System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
                    currentOffsets.put(new TopicPartition(record.topic(), record.partition()),
                            new OffsetAndMetadata(record.offset() + 1, null));
                }
                consumer.commitAsync(currentOffsets, null);
            }
        } catch (WakeupException e) {
            logger.error("Unexpected error", e);
        } finally {
            try {
                consumer.commitSync(currentOffsets);
            } catch (Exception e) {
                throw new RuntimeException(e);
            } finally {
                consumer.close();
            }
        }
    }

    public void seekSpecificOffset() {
        KafkaConsumer<String, String> consumer = create();

        // Assign a specific partition and offset
        TopicPartition partition = new TopicPartition("your-topic", 0);
        consumer.assign(Collections.singleton(partition));
        consumer.seek(partition, 100);

        try {
            while (true) {
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
                for (ConsumerRecord<String, String> record : records) {
                    System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
                }
            }
        } finally {
            consumer.close();
        }
    }

    public void seekOffsetByTime() {
        KafkaConsumer<String, String> consumer = create();
        Long oneHourEarlier = Instant.now().atZone(ZoneId.systemDefault()).minusHours(1).toEpochSecond();
        Map<TopicPartition, Long> partitionLongMap = consumer.assignment().stream().collect(Collectors.toMap(tp -> tp, tp ->oneHourEarlier));
        Map<TopicPartition, OffsetAndTimestamp> offsetMap = consumer.offsetsForTimes(partitionLongMap);
        for (Map.Entry<TopicPartition, OffsetAndTimestamp> entry : offsetMap.entrySet()) {
            consumer.seek(entry.getKey(), entry.getValue().offset());
        }
        try {
            while (true) {
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(100));
                for (ConsumerRecord<String, String> record : records) {
                    System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
                }
            }
        } finally {
            consumer.close();
        }
    }

    private AtomicBoolean exitFlag = new AtomicBoolean(false);

    public void consumeAndExit() {
        KafkaConsumer<String, String> consumer = create();
        addShutdownHook(consumer, Thread.currentThread());
        consumer.subscribe(Arrays.asList("topic1", "topic2"));
        try {
            while (!exitFlag.get()) {
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(10000));
                for (ConsumerRecord<String, String> record : records) {
                    System.out.printf("offset = %d, key = %s, value = %s%n", record.offset(), record.key(), record.value());
                }
                consumer.commitAsync();
            }
            consumer.close();
        } catch (WakeupException e) {
            consumer.commitSync();
        } finally {
            consumer.close();
        }
    }

    public void setExitFlag() {
        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                exitFlag.compareAndSet(false, true);
            }
        });
    }

    public void addShutdownHook(Consumer consumer, Thread mainThread) {
        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                consumer.wakeup();
                try {
                    mainThread.join();
                } catch (InterruptedException e) {
                    logger.error("", e);
                }
            }
        });
    }
}
